/***************************************************************************\
 * Copyright (C) by Keio University
 * UndistorBilinear.cu created in 01 2012.
 * Mail : fdesorbi@hvrl.ics.keio.ac.jp
 *
 * UndistorBilinear.cu is part of the HVRL Engine Library.
 *
 * The HVRL Engine Library is free software; you can redistribute it and/or modify
 * it under the terms of the GNU Lesser General Public License as published by
 * the Free Software Foundation; either version 3 of the License, or
 * (at your option) any later version.
 *
 * The HVRL Engine Library is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Lesser General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 *
\***************************************************************************/


#ifndef RADIAL_DISTORTION_CU__
#define RADIAL_DISTORTION_CU__

#ifdef WIN32
#define NOMINMAX
#include <windows.h>
#endif

#include <hvrl/gpgpu/cuda/Prototypes.hpp>
#include <hvrl/gpgpu/cuda/Tools.cuh>
#include <iostream>


namespace hvrl
{

namespace cuda
{


/**
 * \brief The Kernel computes the undistortion of the image an apply a bilinear interpolation
 *
 * \param map containing precomputed undistort coordinates and parameters for the bilinear interpolation
 * \param input the input RGB color image
 * \param output the output RGBA image
 * \param w the width of the map/input/output
 * \param h the height of the map/input/output
 * \author Francois de Sorbier
 */
__global__ void UndistortBilinear(float *map, // [offsetx|offsety|bilinear param 1|bilinear param 2|bilinear param 3|bilinear param 4]
                                  unsigned char* input, // [R|G|B]
                                  unsigned int* output, // unsigned int [R|G|B|A]
                                  uint w,
                                  uint h)
{
    int x = __umul24(blockIdx.x, blockDim.x) + threadIdx.x;
    int y = __umul24(blockIdx.y, blockDim.y) + threadIdx.y;
    if (x < w && y < h)
    {
        int offsetx = int(map[(y * w + x)*6]);
        int offsety = int(map[(y * w + x)*6+1]);

        if(offsetx >= 0 && offsetx < w-1 && offsety >=0 && offsety < h-1)
        {
            float4 color;
            color.x = color.y = color.z = color.w = 0.0f;

            //   Apply Correction of the radial distortion with bilinear interpolation
            float bi0 = map[(y * w + x)*6+2];
            float bi1 = map[(y * w + x)*6+3];
            float bi2 = map[(y * w + x)*6+4];
            float bi3 = map[(y * w + x)*6+5];

            color.x =
                (input[(offsety*w+offsetx)*3]*bi0+
                 input[(offsety*w+offsetx+1)*3]*bi1+
                 input[((offsety+1)*w+offsetx+1)*3]*bi2+
                 input[((offsety+1)*w+offsetx)*3]*bi3);

            color.y =
                (input[(offsety*w+offsetx)*3+1]*bi0+
                 input[(offsety*w+offsetx+1)*3+1]*bi1+
                 input[((offsety+1)*w+offsetx+1)*3+1]*bi2+
                 input[((offsety+1)*w+offsetx)*3+1]*bi3);

            color.z =
                (input[(offsety*w+offsetx)*3+2]*bi0+
                 input[(offsety*w+offsetx+1)*3+2]*bi1+
                 input[((offsety+1)*w+offsetx+1)*3+2]*bi2+
                 input[((offsety+1)*w+offsetx)*3+2]*bi3);


            output[y * w  + x] = rgbaFloatToInt( color );
        }
    }
}

static float *d_input_map = 0;
static unsigned char *d_input_image = 0;
static int d_width = -1;
static int d_height = -1;
static bool initialized = false;


/**
* \brief Initialisation of the radial distortion
*
* The bilinear interpolation parameters is precomputed.
* \image html BilinearInterpolation.png
*
* \f$f(R_1) \approx \frac{x_2-x}{x_2-x_1} f(Q_{11}) + \frac{x-x_1}{x_2-x_1} f(Q_{21})\f$

* \f$f(R_2) \approx \frac{x_2-x}{x_2-x_1} f(Q_{12}) + \frac{x-x_1}{x_2-x_1} f(Q_{22})\f$

* \f$f(P) \approx \frac{y_2-y}{y_2-y_1} f(R_1) + \frac{y-y_1}{y_2-y_1} f(R_2)\f$
*
* \param map containing precomputed undistort 2D coordinates in float
* \param h_width the width of the input undistort map
* \param h_height the height of the input undistort map
* \return true if everything is fine, otherwise no
* \author Francois de Sorbier
*/
bool initRadialDistortion(
    const float* map, // undistort 2D coordinates map
    const unsigned int& h_width, const unsigned int& h_height // Dimensions of the undistort map
)
{

    // Allocation of the memory on the GPU
    if(d_input_map == 0 && d_input_image == 0)
    {
        checkStatus(cudaMalloc( (void**) &d_input_map, 6*h_width*h_height*sizeof(float)));
        checkStatus(cudaMalloc( (void**) &d_input_image, 3*h_width*h_height*sizeof(unsigned char)));
    }
    else
    {
        return false;
    }

    if((d_width!= int(h_width) || d_height!=int(h_height)))
    {
        d_width=h_width;
        d_height=h_height;
    }

    float* localmap = new float[6*h_width*h_height];

    for(unsigned int i = 0; i < h_height; ++i)
    {
        for(unsigned int j = 0; j < h_width; ++j)
        {
            const float& x = map[(i*h_width+j)*2];
            const float& y = map[(i*h_width+j)*2+1];

            localmap[(i*h_width + j)*6]   = std::floor(x);
            localmap[(i*h_width + j)*6+1] = std::floor(y);

            float x_frac = x - std::floor(x);
            float y_frac = y - std::floor(y);

            // We divide by 255 to avoid extra computation in cuda code
            localmap[(i*h_width + j)*6+2] = (1.0f-x_frac)*(1.0f-y_frac)/255.0f;
            localmap[(i*h_width + j)*6+3] = x_frac*(1.0f-y_frac)/255.0f;
            localmap[(i*h_width + j)*6+4] = x_frac*y_frac/255.0f;
            localmap[(i*h_width + j)*6+5] = (1.0f-x_frac)*y_frac/255.0f;

        }
    }

    // copy the data fron the host to the device
    checkStatus(cudaMemcpy(d_input_map, localmap, 6*d_width*d_height*sizeof(float), cudaMemcpyHostToDevice));

    delete[] localmap;

    initialized = true;
    return true;

}


/**
* \brief Clear the memory allocated on the device
*
* \author Francois de Sorbier
*/
void destroyRadialDistortion(void)
{

    if(d_input_map!=0)
    {
        cudaFree(d_input_map);
    }

    if(d_input_image!=0)
    {
        cudaFree(d_input_image);
    }

    initialized = false;

}

/**
* \brief Update the input color image.
*
* The image is copied to the memory of the device
*
* \author Francois de Sorbier
*/
void updateRadialDistortion(unsigned char* h_input_image)
{
    int size = 3 * d_width * d_height * sizeof(unsigned char);
    checkStatus(cudaMemcpy( d_input_image, h_input_image, size, cudaMemcpyHostToDevice));
}

/**
* \brief Execute the undistort bilinear kernel
*
* The size of the blocks is fixed to 16x16
*
* \author Francois de Sorbier
*/
void runRadialDistortion(unsigned int* d_output)
{
    if(initialized)
    {
        dim3 blockSize(16, 16);
        dim3 gridSize((d_width + blockSize.x - 1) / blockSize.x,
                      (d_height + blockSize.y - 1) / blockSize.y);
        UndistortBilinear<<< gridSize, blockSize>>>(d_input_map, d_input_image, d_output, d_width, d_height);
    }

}

}

}

#endif


